from llm import LLMFactory
from agent import SQLAgent

from app.parameter import PuttyChatParam,KnowledgeChatParam, SQLChatParam
from app.utils import Response
from app.constant import ERRORCODE
from utils import DEBUG,ERROR
from store import VectorStore
from app.service import HubService, DatabaseService
from app.dao import Session
from app.models import File

from prompts import getPuttyCharInput,getKnowledgeCharInput

from embeddings import embeddingPool

class ChatService:
    
    @staticmethod
    def puttyChat(puttyChatParam: PuttyChatParam):
        offical, llm = LLMFactory.get(puttyChatParam.llm)
        llm = llm(puttyChatParam.temperature)
        ERROR(offical)
        ERROR(puttyChatParam)
        # 得到历史聊天数据
        chatInput = getPuttyCharInput(puttyChatParam, offical=offical)
        ERROR(chatInput)
        chatOutput = llm.chat(chatInput)
        return Response(ERRORCODE.RESPONSE_SUCCESS, chatOutput)
    
    @staticmethod
    def knowledgeChat(knowledgeChatParam: KnowledgeChatParam, userId: int):    
        offical, llm = LLMFactory.get(knowledgeChatParam.llm)
        llm = llm(knowledgeChatParam.temperature)
        hubId = knowledgeChatParam.hubId
        # 不基于知识库问答
        if hubId == -1:
            chatInput = getPuttyCharInput(knowledgeChatParam,offical)
            chatOutput = llm.chat(chatInput)
            return Response(ERRORCODE.RESPONSE_SUCCESS, {"chatOutput": chatOutput})
        
        # 进行 hub判定
        res = HubService.judgeHub(userId, hubId)

        # 未找到知识库，返回错误
        if res == None:
            return Response(ERRORCODE.REQUEST_HAVE_NO_PERMISSION)
        res = res.to_dict()
        # 加载向量库
        store = VectorStore(hubId)
        store.load(hubId)
        # 问题查询
        references = store.queryQuestion(knowledgeChatParam.question,embeddingPool.get() ,knowledgeChatParam.referenceCount, knowledgeChatParam.referencePadding)
        DEBUG(references)
        chatInput = getKnowledgeCharInput(knowledgeChatParam, references, offical)
            
        chatOutput = llm.chat(chatInput)
        
        data = {"chatOutput": chatOutput,"references" :references, "hub_name": res["hub_name"]}
        return Response(ERRORCODE.RESPONSE_SUCCESS, data)
    
    @staticmethod
    def sqlChat( sqlChatParam :SQLChatParam, userId: int):
        offical, llm = LLMFactory.get(sqlChatParam.llm)
        llm = llm(sqlChatParam.temperature)
        
        databaseId = sqlChatParam.databaseId
        res = DatabaseService.judgeDatabase(databaseId, userId)
        
        if (res == None):
            return Response(ERROR.REQUEST_HAVE_NO_PERMISSION) 
        
        res = res.to_dict()
        sqlAgent = None
        try:
            sqlAgent = SQLAgent(res['type'], res['username'], res['password'], res['ip'], res["port"], res["database"], llm.llm)
        except Exception as e:
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_SERVER_ERROR, message="Can not establish with the database {}({}:{{}})".format(res['database'], res['ip'], res['port']))
        
        res = sqlAgent.textToSQL(sqlChatParam.question)
        
        return Response(ERRORCODE.RESPONSE_SUCCESS, res)  
        
    
    @staticmethod
    def databaseChat(sqlChatParam, userId: int):
        offical, llm = LLMFactory.get(sqlChatParam.llm)
        llm = llm(sqlChatParam.temperature)
        databaseId = sqlChatParam.databaseId
        res = DatabaseService.judgeDatabase(databaseId, userId)
        
        if (res == None):
            return Response(ERROR.REQUEST_HAVE_NO_PERMISSION) 
        
        res = res.to_dict()
        sqlAgent = None
        try:
            sqlAgent = SQLAgent(res['type'], res['username'], res['password'], res['ip'], res["port"], res["database"], llm.llm)
        except Exception as e:
            ERROR(e)
            return Response(ERRORCODE.RESPONSE_SERVER_ERROR, message="Can not establish with the database {}({}:{{}})".format(res['database'], res['ip'], res['port']))
        
        sql = None
        data = None
        try:
            sql = sqlAgent.textToSQL(sqlChatParam.question)
            data = sqlAgent.sqlQuery(sql)
            data = eval(data)
        except:
            return Response(ERRORCODE.RESPONSE_DATABASE_OPERATION_ERROR, message="SQL execute failed! The SQL AI provide is {}.".format(sql))
        return Response(ERRORCODE.RESPONSE_SUCCESS, {
            "data": data,
            "sql": sql
        })
    
    